"""
Linear classification head for continual learning.

This module implements a multi-head linear classifier approach where each task
has its own dedicated linear classification layer. This architecture ensures
that there's no output interference between tasks, which is a common issue in
continual learning scenarios with shared output spaces.

The multi-head design is particularly effective when task identity is known
during inference (task-aware evaluation), allowing the model to select the
appropriate classifier for each input.
"""

import torch.nn as nn

from continuallearning.models.heads.clfs import CosineClassifier
from continuallearning.models.heads._base_classifier import BaseHead
from continuallearning.registry import HEAD_REGISTRY


@HEAD_REGISTRY.register("linear")
class LinearHead(BaseHead):
    def _create_classifier(self, **kwargs) -> nn.Module:
        """
        Create a task-specific linear classifier module.

        This method creates a linear layer for a specific task with the given
        input features and number of output classes. The classifier is initialized
        using Kaiming normal initialization for weights and zero initialization
        for biases (if enabled).

        Args:
            in_features (int): Number of input features
            num_classes (int): Number of output classes for the task
            *args, **kwargs: Additional arguments (unused in linear implementation)

        Returns:
            nn.Module: Linear classifier module for the task
        """
        classifier = nn.Linear(**self.clf_cfg)
        # Initialize weights using Kaiming normal initialization
        nn.init.xavier_uniform_(classifier.weight)
        if self.bias:
            nn.init.constant_(classifier.bias, 0)

        return classifier


@HEAD_REGISTRY.register("cosine")
class CosineHead(BaseHead):
    def _create_classifier(self, **kwargs) -> nn.Module:
        """
        Create a task-specific cosine classifier module.

        This method creates a cosine similarity-based classifier for a specific task
        with the given input features and number of output classes.

        Args:
            in_features (int): Number of input features
            num_classes (int): Number of output classes for the task
            *args, **kwargs: Additional arguments (unused in cosine implementation)

        Returns:
            nn.Module: Cosine classifier module for the task
        """
        return CosineClassifier(**self.clf_cfg)
